package redisrwmutex

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/go-redsync/redsync/v4"
	"github.com/go-redsync/redsync/v4/redis"
)

// RedisRWMutex 基于RedSync的分布式读写锁
type RedisRWMutex struct {
	rs          *redsync.Redsync
	baseKey     string
	expiry      time.Duration
	retryDelay  time.Duration
	readCounter *redsync.Mutex // 保护读计数器的互斥锁
	writeMutex  *redsync.Mutex // 写锁互斥锁
	instanceMu  sync.Mutex     // 保护本地计数器的互斥锁

	// 本地重入计数
	readCount  int
	writeCount int
}

// NewRedisRWMutex 创建分布式读写锁
func NewRedisRWMutex(pools []redis.Pool, baseKey string) *RedisRWMutex {
	rs := redsync.New(pools...)

	// 默认配置
	defaultExpiry := 10 * time.Second
	defaultRetryDelay := 100 * time.Millisecond

	return &RedisRWMutex{
		rs:          rs,
		baseKey:     baseKey,
		expiry:      defaultExpiry,
		retryDelay:  defaultRetryDelay,
		readCounter: rs.NewMutex(baseKey+":counter", redsync.WithExpiry(defaultExpiry)),
		writeMutex:  rs.NewMutex(baseKey+":write", redsync.WithExpiry(defaultExpiry)),
	}
}

// WithExpiry 设置锁过期时间
func (m *RedisRWMutex) WithExpiry(d time.Duration) *RedisRWMutex {
	m.expiry = d
	m.readCounter = m.rs.NewMutex(m.baseKey+":counter", redsync.WithExpiry(d))
	m.writeMutex = m.rs.NewMutex(m.baseKey+":write", redsync.WithExpiry(d))
	return m
}

// WithRetryDelay 设置重试延迟
func (m *RedisRWMutex) WithRetryDelay(d time.Duration) *RedisRWMutex {
	m.retryDelay = d
	return m
}

// RLock 获取读锁（支持本地重入）
func (m *RedisRWMutex) RLock(ctx context.Context) error {
	m.instanceMu.Lock()
	defer m.instanceMu.Unlock()

	// 本地重入检查
	if m.readCount > 0 {
		m.readCount++
		return nil
	}

	// 如果已经持有写锁，自动获得读权限
	if m.writeCount > 0 {
		m.readCount++
		return nil
	}

	// 检查写锁是否被持有
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			held, err := m.isWriteLockHeld(ctx)
			if err != nil {
				return err
			}
			if !held {
				break
			}
			time.Sleep(m.retryDelay)
		}
	}

	// 获取读计数器锁
	if err := m.readCounter.LockContext(ctx); err != nil {
		return err
	}
	defer m.readCounter.UnlockContext(ctx)

	// 增加读计数
	if err := m.incrementReadCount(ctx); err != nil {
		return err
	}

	m.readCount = 1
	return nil
}

// RUnlock 释放读锁
func (m *RedisRWMutex) RUnlock(ctx context.Context) error {
	m.instanceMu.Lock()
	defer m.instanceMu.Unlock()

	if m.readCount == 0 {
		return fmt.Errorf("read lock not held")
	}

	m.readCount--
	if m.readCount > 0 {
		return nil
	}

	// 完全释放读锁
	if err := m.readCounter.LockContext(ctx); err != nil {
		return err
	}
	defer m.readCounter.UnlockContext(ctx)

	return m.decrementReadCount(ctx)
}

// Lock 获取写锁（支持本地重入）
func (m *RedisRWMutex) Lock(ctx context.Context) error {
	m.instanceMu.Lock()
	defer m.instanceMu.Unlock()

	// 本地重入检查
	if m.writeCount > 0 {
		m.writeCount++
		return nil
	}

	// 获取写锁
	if err := m.writeMutex.LockContext(ctx); err != nil {
		return err
	}

	// 等待所有读锁释放
	for {
		select {
		case <-ctx.Done():
			m.writeMutex.UnlockContext(context.Background())
			return ctx.Err()
		default:
			count, err := m.getReadCount(ctx)
			if err != nil {
				m.writeMutex.UnlockContext(context.Background())
				return err
			}

			if count == 0 {
				m.writeCount = 1
				return nil
			}

			// 续约写锁
			if ok, err := m.writeMutex.ExtendContext(ctx); !ok || err != nil {
				m.writeMutex.UnlockContext(context.Background())
				return fmt.Errorf("failed to extend write lock: %w", err)
			}

			time.Sleep(m.retryDelay)
		}
	}
}

// Unlock 释放写锁
func (m *RedisRWMutex) Unlock(ctx context.Context) error {
	m.instanceMu.Lock()
	defer m.instanceMu.Unlock()

	if m.writeCount == 0 {
		return fmt.Errorf("write lock not held")
	}

	m.writeCount--
	if m.writeCount > 0 {
		return nil
	}

	_, err := m.writeMutex.UnlockContext(ctx)
	return err
}

// ----------------- 辅助方法 -----------------

// 检查写锁是否被持有
func (m *RedisRWMutex) isWriteLockHeld(ctx context.Context) (bool, error) {
	// 检查写锁状态
	if m.writeMutex != nil {
		// 简单的实现：尝试获取写锁，如果能立即获取说明未被持有
		options := []redsync.Option{
			redsync.WithTries(1),
			redsync.WithExpiry(m.expiry),
		}

		tempMutex := m.rs.NewMutex(m.writeMutex.Name(), options...)
		if err := tempMutex.Lock(); err == nil {
			// 成功获取，表示没有写锁
			tempMutex.Unlock()
			return false, nil
		}
	}
	return true, nil
}

// 增加读计数
func (m *RedisRWMutex) incrementReadCount(ctx context.Context) error {
	conn, err := m.rs.Store().Get(context.Background())
	if err != nil {
		return err
	}
	defer conn.Close()

	key := m.baseKey + ":read_count"
	_, err = conn.Incr(ctx, key).Result()
	if err != nil {
		return err
	}

	// 设置过期时间
	_, err = conn.Expire(ctx, key, m.expiry).Result()
	return err
}

// 减少读计数
func (m *RedisRWMutex) decrementReadCount(ctx context.Context) error {
	conn, err := m.rs.Store().Get(context.Background())
	if err != nil {
		return err
	}
	defer conn.Close()

	key := m.baseKey + ":read_count"
	_, err = conn.Decr(ctx, key).Result()
	if err != nil {
		return err
	}

	// 计数器为0时自动删除
	count, err := conn.Get(ctx, key).Int()
	if err != nil || count <= 0 {
		conn.Del(ctx, key)
	}

	return nil
}

// 获取读计数
func (m *RedisRWMutex) getReadCount(ctx context.Context) (int, error) {
	conn, err := m.rs.Store().Get(context.Background())
	if err != nil {
		return 0, err
	}
	defer conn.Close()

	key := m.baseKey + ":read_count"
	val, err := conn.Get(ctx, key).Int()
	if err != nil && err.Error() == "redis: nil" {
		return 0, nil // 键不存在表示计数为0
	}
	return val, err
}
